# Copyright (c) 2025 Huawei Technologies Co. Ltd.
# deepinsight is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#          http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
import os
from typing import Optional
from pymongo import AsyncMongoClient
from pymongo.asynchronous.client_session import AsyncClientSession
from pymongo.asynchronous.collection import AsyncCollection
from deepinsight.config.logger import get_logger

logger = get_logger(__name__)


class MongoDB:
    """MongoDB连接"""
    user = os.getenv('MONGODB_USER', 'admin')
    password = os.getenv('MONGODB_PASSWORD', '')
    host = os.getenv('MONGODB_HOST', 'localhost')
    port = os.getenv('MONGODB_PORT', 27017)
    _client: Optional[AsyncMongoClient] = None

    @classmethod
    def _init_client(cls) -> AsyncMongoClient:
        """初始化MongoDB客户端"""
        if cls._client is None:
            cls._client = AsyncMongoClient(
                f"mongodb://{cls.user}:{cls.password}@{cls.host}:{cls.port}/?directConnection=true&replicaSet=rs0",
                uuidRepresentation="standard"
            )
        return cls._client

    @classmethod
    def get_collection(cls, collection_name: str) -> AsyncCollection:
        """获取MongoDB集合（表）"""
        try:
            client = cls._init_client()
            return client[os.getenv('MONGODB_DATABASE', '')][collection_name]
        except Exception as e:
            logger.exception("[MongoDB] 获取集合 %s 失败", collection_name)
            raise RuntimeError(str(e)) from e

    @classmethod
    async def clear_collection(cls, collection_name: str) -> None:
        """清空MongoDB集合（表）"""
        try:
            client = cls._init_client()
            await client[os.getenv('MONGODB_DATABASE', '')][collection_name].delete_many({})
        except Exception:
            logger.exception("[MongoDB] 清空集合 %s 失败", collection_name)

    @classmethod
    def get_session(cls) -> AsyncClientSession:
        """获取MongoDB会话"""
        client = cls._init_client()
        return client.start_session()
